import os
import cv2
import numpy as np
import tensorflow as tf
from lib.utils.nms_wrapper import nms
from lib.utils.test import im_detect

from lib.nets.vgg16 import vgg16

NETS = {'vgg16': ('vgg16_faster_rcnn_iter_37000.ckpt',)}
DATASETS = {'pascal_voc': ('voc_2007_trainval',), 'pascal_voc_0712': ('voc_2007_trainval+voc_2012_trainval',)}
os.environ["CUDA_VISIBLE_DEVICES"] = "0"


def vis_detections(im, dets, thresh=0.1):
	bac = []
	inds = np.where(dets[:, -1] >= thresh)[0]
	for i in inds:
		bbox = dets[i, :4]
		bac.append(bbox)
	return bac


def demo(sess, net, im):
	bac = [[], [], []]
	scores, boxes = im_detect(sess, net, im)

	CONF_THRESH = 0.25
	NMS_THRESH = 0.2

	for i in range(3):
		bac[i].append(vis_detections(im[:, :, (2, 1, 0)],
		                             dets=np.hstack((boxes, scores[:, 18 + i][:, np.newaxis])).astype(np.float32)[
		                                  nms(np.hstack((boxes, scores[:, 18 + i][:, np.newaxis])).astype(np.float32),
		                                      NMS_THRESH), :],
		                             thresh=CONF_THRESH))

	return np.squeeze(bac)


class rcnn():
	def __init__(self):
		demonet = 'vgg16'
		dataset = 'pascal_voc_0712'
		tfmodel = os.path.join('output', demonet, DATASETS[dataset][0], 'default', NETS[demonet][0])
		if not os.path.isfile(tfmodel + '.meta'):
			print(tfmodel)
			raise IOError(('{:s} not found.\nDid you download the proper networks from '
			               'our server and place them properly?').format(tfmodel + '.meta'))
		tfconfig = tf.ConfigProto(allow_soft_placement=True)
		tfconfig.gpu_options.allow_growth = True
		self.sess = tf.Session(config=tfconfig)
		self.net = vgg16(batch_size=1)

		self.net.create_architecture(self.sess, "TEST", 21, tag='default', anchor_scales=[2, 4, 8])
		saver = tf.train.Saver()
		saver.restore(self.sess, tfmodel)

	def get_box(self, img):
		return demo(self.sess, self.net, img)


if __name__ == '__main__':

	rcnn = rcnn()

	dir = 'G92'
	list = os.listdir(dir)
	print('共计 {} 个文件夹'.format(len(list)))
	for i in range(len(list)):
		name = list[i]
		list2 = os.listdir('{}/{}/'.format(dir, name))
		for j, name2 in enumerate(list2):
			I = cv2.imread('{}/{}/{}'.format(dir, name, name2))
			bac = rcnn.get_box(I)
			total = 0
			total += len(bac[0]) * 3
			total += len(bac[1]) * 2
			total += len(bac[2]) * 1
			print("菌落数量：%s" % total)
			I2 = I.copy()

			for k in range(3):
				if k == 0 and len(bac[k]) > 1:
					for box in bac[k]:
						I2 = cv2.rectangle(I2, (int(box[0]), int(box[1])), (int(box[2]), int(box[3])), (0, 0, 255), 2)
				elif k == 1 and len(bac[k]) > 1:
					for box in bac[k]:
						I2 = cv2.rectangle(I2, (int(box[0]), int(box[1])), (int(box[2]), int(box[3])), (0, 255, 0), 2)
				elif k == 2 and len(bac[k]) > 1:
					for box in bac[k]:
						I2 = cv2.rectangle(I2, (int(box[0]), int(box[1])), (int(box[2]), int(box[3])), (255, 0, 0), 2)
			cv2.imwrite('out_img/{}'.format(name2), I2)
